const express = require("express");
const app = express();
const cors = require('cors');
const { default: axios } = require("axios");

app.use(cors())
app.use(express.json());

app.post('/generate', (req, res) => {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');

  // 通过 axios 获取流数据
  axios.post('http://192.168.51.169:11434/api/generate',
    {
      model: "gemma3:1b",
      prompt: req.body.messages,
    },
    {
      headers: {
        'Content-Type': 'application/json',
      },
      responseType: 'stream',
    })
    .then(response => {
      response.data.on('data', chunk => {
        console.log(chunk.toString())
        res.write(`data: ${chunk.toString()} \n\n`);
      });

      response.data.on('end', () => {
        res.write('data: [DONE] \n\n');
        res.end();
      });

      response.data.on('error', (error) => {
        console.error('Stream error:', error);
        res.status(500).send('Stream error');
      });
    })
    .catch(error => {
      console.error('Error while communicating with Ollama:', error);
      res.status(500).send('Server error');
    });
});

app.listen(8800, (err) => {
  if (err) console.log("服务器启动错误", err)
  console.log("服务器启动成功");
});
